#!/bin/env python

import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import pylab as p
import os
from functions_plot import *
#from plot_items import *
'''
This module contains functions that can be executed to create three different types of plots: 
(1) make_profiles (2) make_scatterplots (3) make_lineplots
When calling each function, it will need two arguments: inputtext.txt and plot_items
'''

#class MyLogFormatter(LogFormatter) :                                                                                                 

#    def __call__(self,x,pos=None) :                                                                                                  
#        if abs(log10(x)) < 2 :                                                                                                       
#            return '$%g$' % (x,)                                                                                                     
#        else :                                                                                                                       
#            return '$10^{%g}$' % (log10(x),)


class MakePlots :
    '''
    The class variables contain all the plot information, so to make new plots, we just need to make new objects of this class, then call the methods within
    Inputtxt format:
    <filepath> <aexp> <id> <color> <legend label> #Currently only using the filepath, not sure what to dow ith the others

    Plot items format:
    <Data type> <X axis> <Y axis>
    '''
    def __init__(self, inputtxt, plot_items) :
        self._inputtxt = inputtxt
        self._plot_items = plot_items
        self._datatype = None
        self._xaxis = None
        self._yaxis = None
        self._xscale = None
        self._yscale = None
        self._xlabelsize = 12
        self._ylabelsize = 12
        self._inputfiles = []

    def parse_inputtxt(self, inputtxt) :
        '''
        Reads in the inputtxt file and adds the files to the list self._inputfiles
        '''
        inputfile = open(inputtxt,'r')

        for line in inputfile: #Going through each line
            lt = line.split() #Splits the line by spaces
            filename = ''
            for item in lt[:1]:
                filename += item #Sets filename to the first part of the line
            print filename
            self._inputfiles.append(filename) #Adds the filename to the list
        
    def parse_plot_items(self, plot_items) :
        '''
        Reads in the plot_items file, which describes what type of data we're plotting, and the X and Y axes
        '''
        plotfile = open(plot_items,'r')

        for line in plotfile: #Goes through each line, but there should only be one line
            lt = line.split()
        self._datatype = lt[0] #First element is the data type
        self._xaxis = lt[1] #Second is X
        self._yaxis = lt[2] #Then Y

    def get_data( self ) :
        '''
        This is where we actually make the plots. Since parse_inputtxt and parse_plot_items have already
        gotten pulled all the required info and stored it into the class variables, this function just
        generates the plots
        '''
        if self._datatype == None : #If for some reason plot_items did not specify the data type
            print 'Data type not specified!'
            return None
	print self._inputfiles #Prints the file list
        for filename in self._inputfiles: #Makes a plot for each file 
            if self._datatype == 'BlackHoleData': #Depending on which data type we are using
                obj = BlackHoleData(filename) #we make obj an instance of that particular class, imported from functions_plot
            elif self._datatype == 'SinkFileData':
                obj = SinkFileData(filename)
            elif self._datatype == 'DarkMatterHaloData':
                obj = DarkMatterHaloData(filename) 
            
            xdata = obj.return_field(self._xaxis) #We can call all the methods of the plot object, so getting the requested data is simple
            ydata = obj.return_field(self._yaxis)
            
            #The below part generates the actual plot. We can add more bells and whistles to the plots as needed
            fig = plt.figure(figsize = (7,7)) 
            subplt = fig.add_subplot(111)
            subplt.scatter(xdata,ydata)
            plt.xlabel(self._xaxis)
            plt.ylabel(self._yaxis)
            n = filename.split('/') #Since the 'filename' we are using is actually the full file path, we split it by '/' so we can pull just the file name in the next line
            plt.title(self._xaxis + ' vs ' +self._yaxis+' ('+n[-1]+')') #n[-1] is the last element of the splitted line, which is just the file name
            plt.savefig(n[-1]+'_'+self._xaxis+'_'+self._yaxis+'.png')
            print 'Saved file '+n[-1]+'_'+self._xaxis+'_'+self._yaxis+'.png' 

    def set_xlabel(self, xlabel, fontsize) :
        pass
    def set_ylabel(self, ylabel, fontsize) :
        pass
    def set_xscale(self, xscale) :
        pass
    def set_ysacle(self, yscale) :
        pass

def make_profiles(inputtxt, plot_items) :
    '''
    Create profiles
    This takes in two arguments: inputtxt and plot_items
    '''
    return

def make_scatterplots(inputtxt, plot_items) :
    '''
    Create scatterplots
    This function makes a new MakePlots object and calls the required methods to make plots to it
    '''
    mp = MakePlots(inputtxt, plot_items)
    mp.parse_inputtxt(inputtxt)
    mp.parse_plot_items(plot_items)
    mp.get_data()

    #subplt.set_yscale('log')
    #subplt.autoscale_view(True,True,True)
'''

    for plot_item in plot_items :
        sp = ScatterPlot( inputtxt, plot_items )  
        pf = p.figure( figsize = (7,7) )
        f = pf.add_subplot(111)############################3
        for dataset in sp.datasets :
            f.scatter( dataset.x, dataset.y, marker=dataset.marker, color=dataset.color, s=100 )
            #f.hold()
        f.set_xlabel( sp.xlabel, fontsize=sp.xlabelsize )
        f.set_ylabel( sp.ylabel, fontsize=sp.ylabelsize )
        f.set_xscale( sp.xscale )
        f.set_yscale( sp.yscale )

        #  Will want a class so I have something that looks like f.set_scatter_plot( sp ) 
        #  which will configure labels, scales, legends, ticks, etc.  Can I add to a pre-existing matplotlib class?
        #  Probably... see how MyLogFormatter() works?  
           
        f.set_title( sp.title )
        pf.savefig( sp.savefig, dpi=100 )
     
    return
'''
def make_lineplots(inputtxt, plot_items) :
    '''
    Create line plots
    This takes in two arguments: inputtxt and plot_items
    '''
'''    parsed_input = parse_plot_input( inputtxt )

    for plot_item in plot_items :
        f = p.figure( figsize=(7,7) )
        axes([0.20,0.15,0.75,0.75])
        plots = []
        for key in parsed_input['keys'] :
            xdata, ydata = plot_out( parsed_input['datadir'][key], 
                                     parsed_input['id'][key],
                                     plot_item
                                     )
            plots.append(f.gca().plot( xdata, ydata, 
                                       parsed_input['linestyle'][key],
                                       markersize=0.3
                                       )
                         )

        f.gca().set_xlabel(plot_item['xlabel'],fontsize='xx-large')
        f.gca().set_ylabel(plot_item['ylabel'],fontsize='xx-large')
        if (plot_item['xlog']) :
            f.gca().set_xscale('log')
            # f.gca().xaxis.set_major_formatter(MyLogFormatter())
        if (plot_item['ylog']) :
            f.gca().set_yscale('log')
            # f.gca().yaxis.set_major_formatter(MyLogFormatter())
        if plot_item['xlim'] != None :
            f.gca().set_xlim(plot_item['xlim'])
        if plot_item['ylim'] != None :
            f.gca().set_ylim(plot_item['ylim'])
        
        if 'legend_label' in parsed_input.keys() :
            f.gca().legend( plots, 
                            [ parsed_input['legend_labels'][key] for key in parsed_input['keys'] ],
                            plot_item['legendposition'],
                            markerscale=4.
                            )
            leg = gca().get_legend()
            legtext, leglines = leg.get_texts(), leg.get_lines()
            setp(legtext,fontsize='large')
            setp(leglines,linewidth=4.)
            leg.drawframe(False)

        f.gca().set_title(plot_item['title'],fontsize='xx-large')

        if not os.path.exists( parsed_input['png_folder'] ) :
            os.mkdir( parsed_input['png_folder'] )
        f.savefig( parsed_input['png_folder']+'/'+plot_item['saveas'],dpi=100 )

    return
'''

